#include "sm2.h"
#include <string.h>
#include <stdio.h>
#include "sm2_rand.h"
#include "sm2_format.h"

extern const SM2_Jacobian_Point SM2_G;
extern const sm2_uint256_t ZERO;
extern const sm2_uint256_t ONE;
extern const sm2_uint256_t SM2_P;
extern const sm2_uint256_t SM2_A;
extern const sm2_uint256_t SM2_B;
extern const sm2_uint256_t SM2_N;
extern const sm2_uint512_t SM2_Ps[13];
extern const sm2_uint512_t SM2_P_sqr;
extern const sm2_uint256_t Gx;
extern const sm2_uint256_t Gy;
extern const sm2_uint256_t SM2_N_MINUS_ONE;


/**
 * @brief   产生一个区间随机大数（256bit）
 * 
 * @details 该函数将根据传入的`low`和`high`两个256bit大数，生成一个随机数`k`落在它们区间
 * 
 * @param[out] k  产生随机数
 * @param[in] low 随机数下界  
 * @param[in] high 随机数上界
 * 
 * @return 无返回值
 * 
 * @note    无
 * @warning 使用时需要确保传入的 `k` 已正确分配内存。
 * @pre     无
 * @post    生成随机数 k 在 [low, high)范围之内
 */
void sm2_rand(sm2_uint256_t k, const sm2_uint256_t low, const sm2_uint256_t high)
{
    sm2_rand_range(k, high);
}

// 顺序：将大数src从高到低拷贝到dist
/**
 * @brief   将大数src根据大端法拷贝到dist
 * 
 * @details 该函数将大数src（小端法表示）从高位到地位依次拷贝到dist内存地址处（大端法），使用uint32指针是为了提升拷贝速度
 * 
 * @param[in, out] dist 目标地址 
 * @param[in] src   源大数地址 
 * 
 * @return 无返回值
 * 
 * @note    无
 * @warning 使用时需要确保传入的 `dist` 已正确分配内存。
 * @pre     无
 * @post    dist地址即后续64字节为大端法排序的大数
 */
void copy_bn(uint32_t *dist, const sm2_uint256_t src)
{
    uint8_t *p = (uint8_t*)(&src[7]), *q = (uint8_t*)dist;
    for(int i=7; i>=0; i--){
        *q++ = *(p+3);
        *q++ = *(p+2);
        *q++ = *(p+1);
        *q++ = *p;
        p -= 4;
    }
}

void copy_reverse(uint8_t *dist, const uint8_t *src, size_t n_bytes)
{
    for(int i=n_bytes-1; i>=0; i--){
        *dist++ = src[i];
    }
}

/**
 * @brief   计算SM2签名流程的Z_A
 * 
 * @details 该函数基于GB/T 32918.2-2016计算用户A杂凑值Z_A = SM3（ENTLA||IDA||a||b||Xg||Yg||Xa||Ya)
 * 
 * @param[in, out]  Z_A 产生的用户Hash Z_A
 * @param[in] pub   公钥
 * @param[in] id    用户可辩别标识ID
 * @param[in] n_bytese 标识长度（单位：字节）
 * 
 * @return 无返回值
 * 
 * @note
 * @warning 使用时需要确保传入的 `pub`, `id` 已正确分配内存，并正确赋值；确保Z_A正确赋值
 * @pre     无
 * @post    得到用户Hash值Z_A
 */
void sm2_compute_z(uint8_t Z_A[32], const SM2_Point *pub, const uint8_t *id, unsigned short n_bytes)
{
    // Z_A=H256(ENTLA||IDA||a||b||xG||yG||xA||yA)
    uint32_t size = 194 + n_bytes;
    uint8_t *concat_m = (uint8_t*)malloc(size); // 194 = 2 + 64 * 3
    // *((uint16_t *)(concat_m)) = n_bytes << 3; //将n_bytes * 8 比特数copy到前两个字节, todo: 顺序
    // uint8_t *p = concat_m;
    uint16_t n_bits = n_bytes << 3;
    copy_reverse(concat_m, (uint8_t*)(&n_bits), 2);
    // p += 2;
    memcpy(concat_m + 2, id, n_bytes);
    // copy_reverse(p, (uint8_t*)id, n_bytes);
    // p += n_bytes;
    // copy_reverse(p, (uint8_t*)SM2_A, 32);
    // p += 32;
    // copy_reverse(p, (uint8_t*)SM2_B, 32);
    // p += 32;
    // copy_reverse(p, (uint8_t*)SM2_G.x, 32);
    // p += 32;
    // copy_reverse(p, (uint8_t*)SM2_G.y, 32);
    // p += 32;
    // copy_reverse(p, (uint8_t*)pub->x, 32);
    // p += 32;
    // copy_reverse(p, (uint8_t*)pub->y, 32);
    uint32_t *p = (uint32_t*)(concat_m + (2 + n_bytes));
    copy_bn(p, SM2_A);
    p+=8;
    copy_bn(p, SM2_B);
    p+=8;
    copy_bn(p, SM2_G.x);;
    p+=8;
    copy_bn(p, SM2_G.y);
    p+=8;
    copy_bn(p, pub->x);
    p+=8;
    copy_bn(p, pub->y);
    // print_bytes(concat_m, size);
    sm3_hash(Z_A, concat_m, size);
    free(concat_m);
}

/**
 * @brief   生成SM2签名(R, S)
 * 
 * @details 该函数基于GB/T 32918.2-2016计算签名
 * 
 * @param[in, out]  sig 签名
 * @param[in] key   用户的密钥对
 * @param[in] Z_A   用户杂凑值Z_A = SM3(EnTL_A||ID_A||a||b||X_g||Y_g||X_a||Y_a)
 * @param[in] message   签名消息
 * @param[in] message_size 消息大小
 * 
 * @return 无返回值
 * 
 * @note    无
 * @warning 使用时要确保`sig`已经分配足够空间，其余参数以及分配足够空间并正确赋值
 * @pre     无
 * @post    得到用户对消息的签名sig
 */
void sm2_signature_generate(const SM2_Key *key, const uint8_t Z_A[32], const uint8_t* message, size_t message_size, SM2_Signature *sig)
{
    uint8_t *_m;
    sm2_uint256_t e, k, r, s, tmp1, tmp2, d_plus_one_inv;
    SM2_Jacobian_Point p, tmp;
    uint8_t e_bytes[32];
    // 预计算 (1+d_A) 关于模n的逆元(1+d_A)^(-1)
    sm2_add(tmp1, ONE, key->private_key);
    sm2_modn_inverse(d_plus_one_inv, tmp1);

    // A1: M' = Z_A||M
    size_t size = 32 + message_size;
    _m = (uint8_t *)malloc(size);
    memcpy(_m, Z_A, 32);
    memcpy(_m + 32, message, message_size);

    // A2: 计算e =Hv(M'), 并将字节串e_bytes转换为整数e
    sm3_hash(e_bytes, _m, size);
    sm2_bytes_to_integer(e, e_bytes);
 
    while(1){
        // A3: 用随机数发生器产生随机数k∈[1, n-1]
        // #define SM2_SET_RANDOM       // 用于验证国家标准文档中的示例
        #ifdef SM2_SET_RANDOM
        char_to_hex(k, "59276E27D506861A16680F3AD9C02DCCEF3CC1FA3CDBE4CE6D54B80DEAC1BC21");
        print_256(k);
        #else
        sm2_rand(k, ONE, SM2_N_MINUS_ONE);
        #endif
        // A4: 计算椭圆曲线点(x1, y1)=[k]G, 并将x1转换为整数
        sm2_jacobian_point_k_mul(&tmp, &SM2_G, k); // (x1, y1)=[k]G, p为素数, x1已经为整数不用转换
        sm2_get_affine_point(&p, &tmp);

        // A5: 计算r = (e+x1) modn，若r = 0或r +k = n则返回A3
        sm2_modn_add(r, e, p.x);
        if(sm2_compare(r, ZERO) == SM2_COMPARE_EQ) continue;
        sm2_mod_add(tmp1, r, k);
        if(sm2_compare(tmp1, SM2_N) == SM2_COMPARE_EQ) continue;

        // A6: 计算s= ((1+dA)^(-1) * (k-r*dA)) mod n，若s=0则返回A3；
        sm2_modn_mul(tmp1, r, key->private_key);
        sm2_modn_sub(tmp2, k, tmp1);
        sm2_modn_mul(s, d_plus_one_inv, tmp2);
        if(sm2_compare(s, ZERO) == SM2_COMPARE_EQ) continue;
        break;
    } 
    // 转换r, s为字节串
    sm2_integer_to_bytes(sig->r, r);
    sm2_integer_to_bytes(sig->s, s);
    free(_m);
}

/**
 * @brief   SM2验证签名
 * 
 * @details 此函数基于GB/T 32918.2-2016对消息签名进行验证，其接受待验证消息、消息签名和用户杂凑值，对消息签名进行验证并返回验证结果
 * 
 * @param[in] public_key 公钥  
 * @param[in] Z_A 用户杂凑值Z_A
 * @param[in] message 待验证消息
 * @param[in] message_size 消息长度
 * @param[in] sig 消息签名
 * 
 * @return 
 *      - 1 签名验证通过
 *      - 0 签名验证未通过
 * 
 * @note    确保消息和消息长度对应
 * @warning 确保传入参数分配足够空间并正确赋值
 * @pre     无
 * @post    得到消息验证结果，消息验证成功则返回1，否则返回0
 */
int sm2_signature_verify(SM2_Point *public_key, const uint8_t Z_A[32], const uint8_t* message, size_t message_size, SM2_Signature *sig)
{
    sm2_uint256_t r, s;
    sm2_bytes_to_integer(r, sig->r);
    sm2_bytes_to_integer(s, sig->s);
    // B1: 检验r’∈[1, n-1] 是否成立
    if(sm2_compare(r, ONE) == SM2_COMPARE_LT || sm2_compare(r, SM2_N_MINUS_ONE) == SM2_COMPARE_GT){
        return 0;
    }
    // B2: 检验s’∈[1, n-1] 是否成立
    if(sm2_compare(s, ONE) == SM2_COMPARE_LT || sm2_compare(s, SM2_N_MINUS_ONE) == SM2_COMPARE_GT){
        return 0;
    }
    // B3: 置M’’=ZA||M’
    uint8_t *_m;
    sm2_uint256_t _e, t;
    uint8_t digest[32];
    size_t size = 32 + message_size;
    _m = (uint8_t *)malloc(size);
    memcpy(_m, Z_A, 32);
    memcpy(_m + 32, message, message_size);
    // B4: 计算e’=Hv(M’’),并将e转换为整数
    sm3_hash(digest, _m, size);
    sm2_bytes_to_integer(_e, digest);
    // B5: 将r’、s’ 的数据类型转换为整数，计算t = (r’+ s’) modn，若t = 0，则验证不通过
    sm2_modn_add(t, r, s);
    free(_m);
    if(sm2_compare(t, ZERO) == SM2_COMPARE_EQ){
        return 0;
    }
    SM2_Jacobian_Point p, tmp1, tmp2;
    // B6: 计算椭圆曲线点(x1’, y1’) = [ s’ ]G + [ t ]PA 
    sm2_jacobian_point_k_mul(&tmp1, &SM2_G, s);
    sm2_set_jacobian_point_from_std(&tmp2, public_key);
    sm2_jacobian_point_k_mul(&p, &tmp2, t);
    sm2_jacobian_point_add(&tmp2, &tmp1, &p);
    sm2_get_affine_point(&p, &tmp2);
    // 计算R = (e’+x1’) modn，检验R = r’是否成立，若成立则验证通过
    sm2_modn_add(t, _e, p.x);
    if(sm2_compare(t, r) != SM2_COMPARE_EQ){
        return 0;
    }
    return 1;
}

int sm2_test()
{
    const char *testhex[] = {
		/* 0 "abc" */
		"616263",
		/* 1 "abcd" 16 times */
		"6162636461626364616263646162636461626364616263646162636461626364"
		"6162636461626364616263646162636461626364616263646162636461626364",
		/* 2 p.57 ZA */
		"0090"
		"414C494345313233405941484F4F2E434F4D"
		"787968B4FA32C3FD2417842E73BBFEFF2F3C848B6831D7E0EC65228B3937E498"
		"63E4C6D3B23B0C849CF84241484BFE48F61D59A5B16BA06E6E12D1DA27C5249A"
		"421DEBD61B62EAB6746434EBC3CC315E32220B3BADD50BDC4C4E6C147FEDD43D"
		"0680512BCBB42C07D47349D2153B70C4E5D7FDFCBFA36EA1A85841B9E46E09A2"
		"0AE4C7798AA0F119471BEE11825BE46202BB79E2A5844495E97C04FF4DF2548A"
		"7C0240F88F1CD4E16352A73C17B7F16F07353E53A176D684A9FE0C6BB798E857",
		/* 3 p.59 ZA */
		"0090"
		"414C494345313233405941484F4F2E434F4D"
		"000000000000000000000000000000000000000000000000000000000000000000"
		"00E78BCD09746C202378A7E72B12BCE00266B9627ECB0B5A25367AD1AD4CC6242B"
		"00CDB9CA7F1E6B0441F658343F4B10297C0EF9B6491082400A62E7A7485735FADD"
		"013DE74DA65951C4D76DC89220D5F7777A611B1C38BAE260B175951DC8060C2B3E"
		"0165961645281A8626607B917F657D7E9382F1EA5CD931F40F6627F357542653B2"
		"01686522130D590FB8DE635D8FCA715CC6BF3D05BEF3F75DA5D543454448166612",
		/* 4 p.72 ZA */
		"0090"
		"414C494345313233405941484F4F2E434F4D"
		"787968B4FA32C3FD2417842E73BBFEFF2F3C848B6831D7E0EC65228B3937E498"
		"63E4C6D3B23B0C849CF84241484BFE48F61D59A5B16BA06E6E12D1DA27C5249A"
		"421DEBD61B62EAB6746434EBC3CC315E32220B3BADD50BDC4C4E6C147FEDD43D"
		"0680512BCBB42C07D47349D2153B70C4E5D7FDFCBFA36EA1A85841B9E46E09A2"
		"3099093BF3C137D8FCBBCDF4A2AE50F3B0F216C3122D79425FE03A45DBFE1655"
		"3DF79E8DAC1CF0ECBAA2F2B49D51A4B387F2EFAF482339086A27A8E05BAED98B",
		/* 5 p.72 ZB */
		"0088"
		"42494C4C343536405941484F4F2E434F4D"
		"787968B4FA32C3FD2417842E73BBFEFF2F3C848B6831D7E0EC65228B3937E498"
		"63E4C6D3B23B0C849CF84241484BFE48F61D59A5B16BA06E6E12D1DA27C5249A"
		"421DEBD61B62EAB6746434EBC3CC315E32220B3BADD50BDC4C4E6C147FEDD43D"
		"0680512BCBB42C07D47349D2153B70C4E5D7FDFCBFA36EA1A85841B9E46E09A2"
		"245493D446C38D8CC0F118374690E7DF633A8A4BFB3329B5ECE604B2B4F37F43"
		"53C0869F4B9E17773DE68FEC45E14904E0DEA45BF6CECF9918C85EA047C60A4C",
		/* 6 p.75 ZA */
		"0090"
		"414C494345313233405941484F4F2E434F4D"
		"000000000000000000000000000000000000000000000000000000000000000000"
		"00E78BCD09746C202378A7E72B12BCE00266B9627ECB0B5A25367AD1AD4CC6242B"
		"00CDB9CA7F1E6B0441F658343F4B10297C0EF9B6491082400A62E7A7485735FADD"
		"013DE74DA65951C4D76DC89220D5F7777A611B1C38BAE260B175951DC8060C2B3E"
		"008E3BDB2E11F9193388F1F901CCC857BF49CFC065FB38B9069CAAE6D5AFC3592F"
		"004555122AAC0075F42E0A8BBD2C0665C789120DF19D77B4E3EE4712F598040415",
		/* 7 p.76 ZB */
		"0088"
		"42494C4C343536405941484F4F2E434F4D"
		"000000000000000000000000000000000000000000000000000000000000000000"
		"00E78BCD09746C202378A7E72B12BCE00266B9627ECB0B5A25367AD1AD4CC6242B"
		"00CDB9CA7F1E6B0441F658343F4B10297C0EF9B6491082400A62E7A7485735FADD"
		"013DE74DA65951C4D76DC89220D5F7777A611B1C38BAE260B175951DC8060C2B3E"
		"0034297DD83AB14D5B393B6712F32B2F2E938D4690B095424B89DA880C52D4A7D9"
		"0199BBF11AC95A0EA34BBD00CA50B93EC24ACB68335D20BA5DCFE3B33BDBD2B62D",
		/* 8 TopsecCA cert ZA */
		"0080"
		"31323334353637383132333435363738"
		"FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFC"
		"28E9FA9E9D9F5E344D5A9E4BCF6509A7F39789F515AB8F92DDBCBD414D940E93"
		"32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7"
		"BC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0"
		"D69C2F1EEC3BFB6B95B30C28085C77B125D77A9C39525D8190768F37D6B205B5"
		"89DCD316BBE7D89A9DC21917F17799E698531F5E6E3E10BD31370B259C3F81C3",
		/* 9 */
		"4D38D2958CA7FD2CFAE3AF04486959CF92C8EF48E8B83A05C112E739D5F181D0"
		"3082020CA003020102020900"
		"AF28725D98D33143300C06082A811CCF"
		"550183750500307D310B300906035504"
		"060C02636E310B300906035504080C02"
		"626A310B300906035504070C02626A31"
		"0F300D060355040A0C06746F70736563"
		"310F300D060355040B0C06746F707365"
		"633111300F06035504030C08546F7073"
		"65634341311F301D06092A864886F70D"
		"0109010C10626A40746F707365632E63"
		"6F6D2E636E301E170D31323036323430"
		"37353433395A170D3332303632303037"
		"353433395A307D310B30090603550406"
		"0C02636E310B300906035504080C0262"
		"6A310B300906035504070C02626A310F"
		"300D060355040A0C06746F7073656331"
		"0F300D060355040B0C06746F70736563"
		"3111300F06035504030C08546F707365"
		"634341311F301D06092A864886F70D01"
		"09010C10626A40746F707365632E636F"
		"6D2E636E3059301306072A8648CE3D02"
		"0106082A811CCF5501822D03420004D6"
		"9C2F1EEC3BFB6B95B30C28085C77B125"
		"D77A9C39525D8190768F37D6B205B589"
		"DCD316BBE7D89A9DC21917F17799E698"
		"531F5E6E3E10BD31370B259C3F81C3A3"
		"733071300F0603551D130101FF040530"
		"030101FF301D0603551D0E041604148E"
		"5D90347858BAAAD870D8BDFBA6A85E7B"
		"563B64301F0603551D23041830168014"
		"8E5D90347858BAAAD870D8BDFBA6A85E"
		"7B563B64300B0603551D0F0404030201"
		"06301106096086480186F84201010404"
		"03020057",
	};

	const char *dgsthex[] = {
		"66c7f0f462eeedd9d1f2d46bdc10e4e24167c4875cf2f7a2297da02b8f4ba8e0",
		"debe9ff92275b8a138604889c18e5a4d6fdb70e5387e5765293dcba39c0c5732",
		"F4A38489E32B45B6F876E3AC2168CA392362DC8F23459C1D1146FC3DBFB7BC9A",
		"26352AF82EC19F207BBC6F9474E11E90CE0F7DDACE03B27F801817E897A81FD5",
		"E4D1D0C3CA4C7F11BC8FF8CB3F4C02A78F108FA098E51A668487240F75E20F31",
		"6B4B6D0E276691BD4A11BF72F4FB501AE309FDACB72FA6CC336E6656119ABD67",
		"329c2f6030cc7e0ca3af6c97b76243ca250338ad3d3dc3a8b322d1cfdf98c2b7",
		/*"ECF0080215977B2E5D6D61B98A99442F03E8803DC39E349F8DCA5621A9ACDF2B",*/
		"557BAD30E183559AEEC3B2256E1C7C11F870D22B165D015ACF9465B09B87B527",
		"4D38D2958CA7FD2CFAE3AF04486959CF92C8EF48E8B83A05C112E739D5F181D0",
		"C3B02E500A8B60B77DEDCF6F4C11BEF8D56E5CDE708C72065654FD7B2167915A",
	};

    uint8_t result[32];
    uint8_t message[1024];
    uint8_t answer[32];
    sm2_uint256_t num = {0};
    for(int i=0; i<10; i++){
        convert_hex_message_to_bytes(message, testhex[i]);
        sm3_hash(result, message, strlen(testhex[i]) / 2);
        convert_hex_message_to_bytes(answer, dgsthex[i]);
        if(sm2_bytes_compare(result, answer) == 0){
            printf("the %d message is correct\n", i+1);
        }else{
            printf("the %d message is wrong\n", i+1);
        }
        for(int j = 0; j < 32; j++){
            printf("%02X", result[j]);
        }
        putchar('\n');
        // break;
    }
    return 0;
}

void print_bytes(const uint8_t *s, size_t n)
{
    for(int i=0; i<n; i++){
        printf("%02X", s[i]);
    }
    putchar('\n');
}


void test()
{
    SM2_Key key;
    char *pri_key = "3945208F7B2144B13F36E38AC6D39F95889393692860B51A42FB81EF4DF7C5B8";
    sm2_uint256_t num = {0};
    char_to_hex(num, pri_key);
    print_256(num);
    char_to_hex(key.private_key, pri_key);
    SM2_Jacobian_Point pub;
    // sm2_key_generate(&key);
    sm2_jacobian_point_k_mul(&pub, &SM2_G, key.private_key);
    sm2_get_std_point(&(key.public_key), &pub);
    sm2_key_print(stdout, 0, 4, "privateKey", &key);
    print_256(key.public_key.x);
    print_256(key.public_key.y);
    char *message = "message digest";
    char *id = "1234567812345678";

    uint8_t Z_A[32] = {0};
    SM2_Signature sig;
    sm2_compute_z(Z_A, &key.public_key, id, strlen(id));
    printf("Z_A: ");
    print_bytes(Z_A, 32);
    printf("\n\n");
    sm2_signature_generate(&key, Z_A, message, strlen(message), &sig);
    sm2_uint256_t r, s;
    sm2_bytes_to_integer(r, sig.r);
    sm2_bytes_to_integer(s, sig.s);
    print_256(r);
    print_256(s);
    // Z_A[0]='f';
    // sig.s[0] = 'a';
    int ret = sm2_signature_verify(&key.public_key, Z_A, message, strlen(message), &sig);
    if(ret){
        printf("success\n");
    }else{
        printf("failed\n");
    }
}

// int main(void)
// {
//     test();
//     return 0;
// }
